package io.kwdb.sdk;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;
import java.util.function.Consumer;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * Batch processor for accumulating items and submitting them in batches. Supports configurable
 * batch size, linger time, and backpressure strategies.
 *
 * @param <T> Type of items to batch
 */
public final class Batcher<T> implements AutoCloseable {

  private static final Logger logger = Logger.getLogger(Batcher.class.getName());

  private final BlockingDeque<T> queue;
  private final int maxBatchSize;
  private final int lingerMs;
  private final Config.Backpressure backpressure;
  private final Consumer<List<T>> sink;
  private final Thread workerThread;
  private final CountDownLatch shutdownLatch = new CountDownLatch(1);
  private final Object flushSignal = new Object();

  // Statistics
  private final AtomicLong droppedCount = new AtomicLong(0);
  private final AtomicLong submittedBatches = new AtomicLong(0);
  private final AtomicLong submittedItems = new AtomicLong(0);

  private volatile boolean closed = false;
  private volatile boolean flushRequested = false;

  /**
   * Create a new batcher with specified configuration.
   *
   * @param maxBatchSize Maximum number of items per batch
   * @param lingerMs Maximum time to wait before submitting a partial batch
   * @param maxQueueSize Maximum queue size
   * @param backpressure Backpressure strategy when queue is full
   * @param sink Consumer to process batches
   */
  public Batcher(
      int maxBatchSize,
      int lingerMs,
      int maxQueueSize,
      Config.Backpressure backpressure,
      Consumer<List<T>> sink) {
    if (maxBatchSize <= 0) {
      throw new IllegalArgumentException("maxBatchSize must be positive");
    }
    if (maxQueueSize <= 0) {
      throw new IllegalArgumentException("maxQueueSize must be positive");
    }
    if (sink == null) {
      throw new IllegalArgumentException("sink cannot be null");
    }

    this.maxBatchSize = maxBatchSize;
    this.lingerMs = lingerMs;
    this.backpressure = backpressure != null ? backpressure : Config.Backpressure.BLOCK;
    this.sink = sink;
    this.queue = new LinkedBlockingDeque<>(maxQueueSize);

    // Start worker thread
    this.workerThread = new Thread(this::run, "Batcher-Worker");
    this.workerThread.setDaemon(true);
    this.workerThread.start();

    logger.log(
        Level.INFO,
        "Batcher started: maxBatchSize={0}, lingerMs={1}, maxQueueSize={2}, backpressure={3}",
        new Object[] {maxBatchSize, lingerMs, maxQueueSize, backpressure});
  }

  /**
   * Offer an item to the batcher.
   *
   * @param item Item to add
   * @param timeoutMs Maximum time to wait in BLOCK mode
   * @return true if item was accepted, false otherwise
   */
  public boolean offer(T item, long timeoutMs) {
    if (closed) {
      return false;
    }

    try {
      if (backpressure == Config.Backpressure.BLOCK) {
        // Block until space is available or timeout
        return queue.offerLast(item, timeoutMs, TimeUnit.MILLISECONDS);
      } else { // DROP_OLDEST
        if (!queue.offerLast(item)) {
          // Queue is full, drop oldest
          T dropped = queue.pollFirst();
          if (dropped != null) {
            droppedCount.incrementAndGet();
            logger.log(
                Level.WARNING,
                "Dropped oldest item due to backpressure, total dropped: {0}",
                droppedCount.get());
          }
          return queue.offerLast(item);
        }
        return true;
      }
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      return false;
    }
  }

  /** Flush any pending items immediately. Blocks until the current batch is submitted. */
  public void flush() {
    if (closed) {
      return;
    }

    synchronized (flushSignal) {
      flushRequested = true;
      flushSignal.notifyAll();
    }

    // Wait for flush to complete
    // Need to wait at least lingerMs + buffer for worker to process and submit batch
    try {
      long waitTime = Math.max(100, lingerMs + 100); // lingerMs + 100ms buffer
      Thread.sleep(waitTime);
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
    }
  }

  @Override
  public void close() {
    if (closed) {
      return;
    }

    logger.log(Level.INFO, "Closing batcher...");

    // Flush remaining items BEFORE setting closed flag
    // This ensures pending writes can still complete
    flush();

    // Now set closed to prevent new items
    closed = true;

    // Wait for worker thread to finish naturally (don't interrupt immediately)
    // Worker will exit when it sees closed=true and queue is empty
    try {
      if (!shutdownLatch.await(5, TimeUnit.SECONDS)) {
        logger.log(Level.WARNING, "Batcher worker thread did not terminate in time");
        // Only interrupt as last resort if worker hangs
        workerThread.interrupt();
      }
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      logger.log(Level.WARNING, "Interrupted while waiting for batcher shutdown");
    }

    // Log statistics
    logger.log(
        Level.INFO,
        "Batcher closed. Stats: batches={0}, items={1}, dropped={2}",
        new Object[] {submittedBatches.get(), submittedItems.get(), droppedCount.get()});
  }

  /** Worker thread main loop. */
  private void run() {
    List<T> batch = new ArrayList<>(maxBatchSize);
    long lastSubmitTime = System.currentTimeMillis();

    try {
      while (!closed) {
        try {
          // Calculate timeout based on linger time
          long now = System.currentTimeMillis();
          long timeSinceLastSubmit = now - lastSubmitTime;
          long timeout = lingerMs > 0 ? Math.max(1, lingerMs - timeSinceLastSubmit) : 1000;

          // Check for flush request
          boolean shouldFlush = false;
          synchronized (flushSignal) {
            if (flushRequested) {
              shouldFlush = true;
              flushRequested = false;
            }
          }

          // Poll with timeout
          T item = queue.poll(timeout, TimeUnit.MILLISECONDS);

          if (item != null) {
            batch.add(item);

            // Drain more items up to batch size
            while (batch.size() < maxBatchSize) {
              T next = queue.poll();
              if (next == null) break;
              batch.add(next);
            }
          }

          // Submit batch if conditions are met
          boolean shouldSubmit =
              !batch.isEmpty()
                  && (batch.size() >= maxBatchSize
                      || // Batch is full
                      shouldFlush
                      || // Flush requested
                      (lingerMs > 0 && timeSinceLastSubmit >= lingerMs) // Linger timeout
                  );

          if (shouldSubmit) {
            submitBatch(batch);
            batch = new ArrayList<>(maxBatchSize);
            lastSubmitTime = System.currentTimeMillis();
          }

        } catch (InterruptedException e) {
          // Expected on shutdown
          if (!closed) {
            logger.log(Level.WARNING, "Worker thread interrupted unexpectedly");
          }
          break;
        } catch (Exception e) {
          logger.log(Level.SEVERE, "Error in batcher worker thread", e);
          // Continue running unless closed
        }
      }

      // Submit any remaining items on shutdown
      if (!batch.isEmpty()) {
        submitBatch(batch);
      }

    } finally {
      shutdownLatch.countDown();
    }
  }

  /** Submit a batch to the sink. */
  private void submitBatch(List<T> batch) {
    if (batch.isEmpty()) {
      return;
    }

    try {
      long startTime = System.currentTimeMillis();
      sink.accept(batch);
      long elapsed = System.currentTimeMillis() - startTime;

      submittedBatches.incrementAndGet();
      submittedItems.addAndGet(batch.size());

      logger.log(
          Level.FINE,
          "Submitted batch: size={0}, elapsed={1}ms",
          new Object[] {batch.size(), elapsed});

    } catch (Exception e) {
      logger.log(Level.SEVERE, "Error submitting batch of " + batch.size() + " items", e);
      // Note: In production, might want to retry or handle differently
    }
  }

  /** Get statistics about the batcher. */
  public BatcherStats getStats() {
    return new BatcherStats(
        submittedBatches.get(), submittedItems.get(), droppedCount.get(), queue.size());
  }

  /** Statistics about batcher performance. */
  public static class BatcherStats {
    public final long submittedBatches;
    public final long submittedItems;
    public final long droppedItems;
    public final int currentQueueSize;

    public BatcherStats(
        long submittedBatches, long submittedItems, long droppedItems, int currentQueueSize) {
      this.submittedBatches = submittedBatches;
      this.submittedItems = submittedItems;
      this.droppedItems = droppedItems;
      this.currentQueueSize = currentQueueSize;
    }

    @Override
    public String toString() {
      return String.format(
          "BatcherStats{batches=%d, items=%d, dropped=%d, queueSize=%d}",
          submittedBatches, submittedItems, droppedItems, currentQueueSize);
    }
  }
}
